#include "yolov3.h"
#include <QDebug>

yolov3::yolov3(QWidget *parent) : QWidget(parent)
{
    modelWeights=new QString("res\\better.weights");
    modelCfg=new QString("res\\better.cfg");
    classesFile=new QString("res\\helmet.names");
    classes=new vector<string>;
    setClasses();
    net=nullptr;
    loadNet();
}
void yolov3::modifyNet(QString cfg,QString name,QString weight,double confidence)
{
    setmodelCfg(cfg);
    setmodelWeights(weight);
    setClassesFile(name);
    setConfidence(confidence);
    setClasses();
    //    qDebug()<<cfg;
    //    qDebug()<<name;
    //    qDebug()<<weight;
    loadNet();
}
float yolov3::getConfidence()
{
    return this->confThreshold;
}
QString yolov3::getmodelWeights()
{
    return *this->modelWeights;
}
void yolov3::setmodelWeights(QString modelWeights)
{
    *this->modelWeights=modelWeights;
}
QString yolov3::getmodelCfg()
{
    return *this->modelCfg;
}
void yolov3::setmodelCfg(QString ModeCfg)
{
    *this->modelCfg=ModeCfg;
}
QString yolov3::getClassesFile()
{
    return *this->classesFile;
}
void yolov3::setClassesFile(QString classesFile)
{
    *this->classesFile=classesFile;
}

void yolov3::setClasses()
{
    ifstream ifs(classesFile->toStdString());
    string line;
    while (getline(ifs, line)) classes->push_back(line);
}
void yolov3::setConfidence(double confidence)
{
    this->confThreshold=confidence;
    qDebug()<<confThreshold;
}
void yolov3::postprocess(cv::Mat& frame, const vector<cv::Mat>& outs)
{
    vector<int> classIds;
    vector<float> confidences;
    vector<cv::Rect> boxes;

    for (size_t i = 0; i < outs.size(); ++i)
    {
        float* data = (float*)outs[i].data;
        for (int j = 0; j < outs[i].rows; ++j, data += outs[i].cols)
        {
            cv::Mat scores = outs[i].row(j).colRange(5, outs[i].cols);
            cv::Point classIdPoint;
            double confidence;

            minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);
            if (confidence > confThreshold)
            {
                int centerX = (int)(data[0] * frame.cols);
                int centerY = (int)(data[1] * frame.rows);
                int width = (int)(data[2] * frame.cols);
                int height = (int)(data[3] * frame.rows);
                int left = centerX - width / 2;
                int top = centerY - height / 2;

                classIds.push_back(classIdPoint.x);
                confidences.push_back((float)confidence);
                boxes.push_back(cv::Rect(left, top, width, height));
            }
        }
    }

    vector<int> indices;
    cv::dnn::NMSBoxes(boxes, confidences, confThreshold, nmsThreshold, indices);
    for (size_t i = 0; i < indices.size(); ++i)
    {
        int idx = indices[i];
        cv::Rect box = boxes[idx];
        drawPred(classIds[idx], confidences[idx], box.x, box.y,
                 box.x + box.width, box.y + box.height, frame);
    }
}

void yolov3::drawPred(int classId, float conf, int left, int top, int right, int bottom, cv::Mat& frame)
{
    string label = cv::format("%.2f", conf);
    cv::Scalar color;
    if (!classes->empty())
    {
        CV_Assert(classId < (int)classes->size());
        string sign;
        if(classes->at(classId)=="person")
        {
            sign="NO!!";
            color = cv::Scalar(0,0,255);
            emit violation(frame);
        }
        else
        {
            sign="YES";
            color = cv::Scalar(0,255,0);
        }

        label = sign + ":" + label;
    }
    rectangle(frame, cv::Point(left, top), cv::Point(right, bottom), color, 2);
    int baseLine;
    cv::Size labelSize = getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
    top = max(top, labelSize.height);
    //    rectangle(frame, cv::Point(left, top - round(1.5*labelSize.height)), cv::Point(left + round(1.5*labelSize.width), top + baseLine), cv::Scalar(255, 255, 255,0), cv::FILLED);
    putText(frame, label, cv::Point(left, top), cv::FONT_HERSHEY_SIMPLEX, 0.75, color,2,cv::LINE_AA);
}

vector<cv::String> yolov3::getOutputsNames(const cv::dnn::Net& net)
{
    static vector<cv::String> names;
    if (names.empty())
    {
        vector<int> outLayers = net.getUnconnectedOutLayers();

        vector<cv::String> layersNames = net.getLayerNames();

        names.resize(outLayers.size());
        for (size_t i = 0; i < outLayers.size(); ++i)
            names[i] = layersNames[outLayers[i] - 1];
    }
    return names;
}

void yolov3::loadNet()
{
    if(net!=nullptr)
        delete net;

    net = new cv::dnn::Net(cv::dnn::readNetFromDarknet(modelCfg->toStdString(), modelWeights->toStdString()));

    //CPU
//    net->setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
//    net->setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
    //GPU
    net->setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
    net->setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
}

cv::Mat yolov3::detect(cv::Mat *frame)
{
    cv::Mat blob;
    cv::dnn::blobFromImage(*frame, blob, 1/255.0, cv::Size(inpWidth, inpHeight), cv::Scalar(0,0,0), true, false);

    net->setInput(blob);

    vector<cv::Mat> outs;
    net->forward(outs, getOutputsNames(*net));

    postprocess(*frame, outs);

    //    vector<double> layersTimes;
    //    double freq = cv::getTickFrequency() / 1000;
    //    double t = net->getPerfProfile(layersTimes) / freq;
    //    string label = cv::format("Inference time for a frame : %.2f ms", t);
    //    putText(*frame, label, cv::Point(0, 15), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 255));

    //    cv::Mat detectedFrame;
    //    frame->convertTo(detectedFrame, CV_8U);
    return *frame;
}

yolov3::~yolov3()
{
    delete modelWeights;
    delete modelCfg;
    delete classesFile;
    delete classes;
}
